/*=============================================================================
*
*   Copyright (C) 2020 All rights reserved.
*
*   Filename:		vtu_writer.hpp
*
*   Author: Wang Zhecheng - wangzhecheng@yeah.net
*
*   Date: 2020-03-13 12:30
*
*   Last Editors: Wang Zhecheng - wangzhecheng@yeah.net
*
*   Last modified:	2020-03-30 21:10
*
*   Description:
*
=============================================================================*/
#ifndef VTK_WRITER_H
#define VTK_WRITER_H

#include "../model/model.hpp"
#include "data_writer.hpp"

namespace geoxfem {
namespace io {

struct vtu {};

struct vtu_node_info { /*{{{*/
  template <typename Mesh, typename NodeFunc>
  static boost::property_tree::ptree apply(
    const Mesh&                                             mesh,
    const NodeFunc&                                         func,
    const std::vector<std::pair<std::string, std::string>>& attr_list = {}) {
    std::string              content = "\n";
    std::vector<std::string> content_list(mesh.getNumberOfVertices());
    std::string              default_content = func(**mesh.getShapeBegin(0));
    for(auto& vertex : shapesIn(mesh, 0)) {
      int id = vertex->getIndex();
      if(id >= 0) {
        if(id >= content_list.size())
          content_list.resize(id);
        content_list.at(id) = func(*vertex);
      }
    }
    for(auto& v : content_list) {
      if(v.empty())
        content.append(default_content);
      content.append(v);
    }
    boost::property_tree::ptree DataArray;
    DataArray.put_value(content);
    for(auto& v : attr_list)
      DataArray.add("<xmlattr>." + v.first, v.second);
    return DataArray;
  }
};                     /*}}}*/
struct vtu_cell_info { /*{{{*/
  template <typename Mesh, typename CellFunc>
  static boost::property_tree::ptree apply(
    const Mesh&                                             mesh,
    const CellFunc&                                         func,
    const std::vector<std::pair<std::string, std::string>>& attr_list = {}) {
    size_t      topo_dim = traits::get_topology_dimension<Mesh>::apply(mesh);
    std::string content  = "\n";
    std::vector<std::string> content_list(mesh.getNumberOfCells());
    for(auto& cell : shapesIn(mesh, topo_dim)) {
      int id = cell->getIndex();
      if(id >= 0) {
        if(id >= content_list.size())
          content_list.resize(id + 1);
        content_list.at(id) = func(*cell);
      }
    }
    std::string default_content = func(**mesh.getShapeBegin(topo_dim));
    for(auto& v : content_list) {
      if(v.empty())
        content.append(default_content);
      content.append(v + "\n");
    }
    boost::property_tree::ptree DataArray;
    DataArray.put_value(content);
    for(auto& v : attr_list)
      DataArray.add("<xmlattr>." + v.first, v.second);
    return DataArray;
  }
};                      /*}}}*/
struct vtu_gauss_info { /*{{{*/
  template <typename Mesh, typename GaussFunc>
  static boost::property_tree::ptree apply(
    const Mesh&                                             mesh,
    const GaussFunc&                                        func,
    const std::vector<std::pair<std::string, std::string>>& attr_list = {}) {
    size_t      topo_dim = traits::get_topology_dimension<Mesh>::apply(mesh);
    std::string content  = "\n";
    for(auto& cell : shapesIn(mesh, topo_dim))
      for(auto& gp : gpsIn(*cell))
        content.append(func(gp));
    boost::property_tree::ptree DataArray;
    DataArray.put_value(content);
    for(auto& v : attr_list)
      DataArray.add("<xmlattr>." + v.first, v.second);
    return DataArray;
  }
}; /*}}}*/

template <typename Model>
struct DataWriter<Model, vtu> : util::ObserverInterface<Model> {
  using ptree = boost::property_tree::ptree;

  void init(const std::filesystem::path& path, /*{{{*/
            const std::filesystem::path& name) {
    m_path      = path;
    m_proj_name = name;
    m_path.append(name.string());
    std::filesystem::create_directory(m_path);
    clearDirectory(m_path);
    m_count = 0;
    m_pvd.put("VTKFile.<xmlattr>.type", "Collection");
    m_pvd.put("VTKFile.<xmlattr>.version", "0.1");
    m_pvd.put("VTKFile.<xmlattr>.byte_order", "LittleEndian");           //
    m_pvd.put("VTKFile.<xmlattr>.compressor", "vtkZLibDataCompressor");  //
    m_pvd.add_child("VTKFile.Collection", ptree{});

    m_count2 = 0;
    m_gauss_pvd.put("VTKFile.<xmlattr>.type", "Collection");
    m_gauss_pvd.put("VTKFile.<xmlattr>.version", "0.1");
    m_gauss_pvd.put("VTKFile.<xmlattr>.byte_order", "LittleEndian");  //
    m_gauss_pvd.put("VTKFile.<xmlattr>.compressor",
                    "vtkZLibDataCompressor");  //
    m_gauss_pvd.add_child("VTKFile.Collection", ptree{});
    return;
  } /*}}}*/

  void update(const Model* model) final override { /*{{{*/
    auto root = create_vtu();
    auto mesh = *model->getMesh();

    size_t num_vertex = mesh.getNumberOfShapes(0);
    root.add("VTKFile.UnstructuredGrid.Piece.<xmlattr>.NumberOfPoints",
             num_vertex);
    size_t num_cell = mesh.getNumberOfShapes(mesh.getTopologyDimension());
    root.add("VTKFile.UnstructuredGrid.Piece.<xmlattr>.NumberOfCells",
             num_cell);

    ptree& Points = root.get_child("VTKFile.UnstructuredGrid.Piece.Points");
    ptree& Cells  = root.get_child("VTKFile.UnstructuredGrid.Piece.Cells");
    ////////////////////////////////////////////////////////////////
    // populate Points

    ptree coords = io::vtu_node_info::apply(
      mesh, io::vtu_vertex_coord{},
      {{"type", "Float32"}, {"NumberOfComponents", "3"}});
    Points.add_child("DataArray", coords);
    ////////////////////////////////////////////////////////////////
    // populate cell

    /* reset count of cell offset at beginning of each output. An ugly
     * workaround, I known... There might be a more elegant way */
    io::vtu_cell_offset{}.operator()(typename Model::Shape{});
    ptree                 offsets = io::vtu_cell_info::apply(
      mesh, io::vtu_cell_offset{}, {{"type", "Int32"}, {"Name", "offsets"}});
    Cells.add_child("DataArray", offsets);

    ptree connectivity = io::vtu_cell_info::apply(
      mesh, [](const auto& v) { return v.getVTKPointOrder(); },
      {{"type", "Int32"}, {"Name", "connectivity"}});
    Cells.add_child("DataArray", connectivity);

    ptree types = io::vtu_cell_info::apply(
      mesh, io::vtu_cell_type{}, {{"type", "UInt8"}, {"Name", "types"}});
    Cells.add_child("DataArray", types);
    ////////////////////////////////////////////////////////////////

    write_ptree<typename Model::CouplePolicy>::apply(mesh, root);

    write_vtu(root, model->getTime() + model->getIter() * 1e-4);
    ////////////////////////////////////////////////////////////////
    ////////////////////////////////////////////////////////////////
    ////////////////////////////////////////////////////////////////

    auto   root2 = create_vtu();
    ptree& gauss_Points =
      root2.get_child("VTKFile.UnstructuredGrid.Piece.Points");
    ptree& gauss_data =
      root2.get_child("VTKFile.UnstructuredGrid.Piece.PointData");

    size_t count = 0;
    for(auto cell : shapesIn(mesh, mesh.getTopologyDimension()))
      count += cell->getNumberOfGaussPoints();
    root2.put("VTKFile.UnstructuredGrid.Piece.<xmlattr>.NumberOfPoints", count);
    root2.put("VTKFile.UnstructuredGrid.Piece.<xmlattr>.NumberOfCells", 0);

    ptree gauss_offset, gauss_connectivity, gauss_type;
    gauss_offset.put("<xmlattr>.type", "Int32");
    gauss_offset.put("<xmlattr>.Name", "offset");
    gauss_connectivity.put("<xmlattr>.type", "Int32");
    gauss_connectivity.put("<xmlattr>.Name", "connectivity");
    gauss_type.put("<xmlattr>.type", "Int32");
    gauss_type.put("<xmlattr>.Name", "types");

    root2.add_child("VTKFile.UnstructuredGrid.Piece.Cells.DataArray",
                    gauss_offset);
    root2.add_child("VTKFile.UnstructuredGrid.Piece.Cells.DataArray",
                    gauss_connectivity);
    root2.add_child("VTKFile.UnstructuredGrid.Piece.Cells.DataArray",
                    gauss_type);

    ptree gauss_coord = io::vtu_gauss_info::apply(
      mesh, io::vtu_gauss_coord{},
      {{"type", "Float32"}, {"NumberOfComponents", "3"}});
    gauss_Points.add_child("DataArray", gauss_coord);

    ptree gauss_sigma = io::vtu_gauss_info::apply(
      mesh, io::vtu_gauss_sigma{},
      {{"type", "Float32"}, {"NumberOfComponents", "9"}, {"Name", "stress"}});
    gauss_data.add_child("DataArray", gauss_sigma);

    ptree gauss_epsilon = io::vtu_gauss_info::apply(
      mesh, io::vtu_gauss_epsilon{},
      {{"type", "Float32"}, {"NumberOfComponents", "9"}, {"Name", "strain"}});
    gauss_data.add_child("DataArray", gauss_epsilon);

    write_gauss_vtu(root2, model->getTime() + model->getIter() * 1e-4);

    return;
  } /*}}}*/

private:
  void clearDirectory(const std::filesystem::path& path) { /*{{{*/
    if(path.empty())
      return;
    for(auto& p : std::filesystem::recursive_directory_iterator(path))
      if(p.path().extension() == ".vtu")
        std::filesystem::remove(p);
    return;
  }                                                            /*}}}*/
  void write_gauss_vtu(const ptree& root, const _real& time) { /*{{{*/
    std::filesystem::path vtu_full_path = m_path;
    std::filesystem::path file_name =
      m_proj_name.string() + "_gauss_" + std::to_string(++m_count2) + ".vtu";
    vtu_full_path.append(file_name.string());
    boost::property_tree::write_xml(
      vtu_full_path, root, std::locale(),
      boost::property_tree::xml_writer_make_settings<std::string>(' ', 2));
    ////////////////////////////////////////////////////////////////

    std::filesystem::path pvd_full_path = m_path;
    pvd_full_path.append(m_proj_name.string() + "_gauss_" + ".pvd");
    ptree& collection = m_gauss_pvd.get_child("VTKFile.Collection");
    ptree  DataSet;

    DataSet.put("<xmlattr>.timestep", time);
    DataSet.put("<xmlattr>.group", "");
    DataSet.put("<xmlattr>.part", "0");
    DataSet.put("<xmlattr>.file", file_name.string());
    collection.add_child("DataSet", DataSet);

    boost::property_tree::write_xml(
      pvd_full_path, m_gauss_pvd, std::locale(),
      boost::property_tree::xml_writer_make_settings<std::string>(' ', 2));
    return;
  }                                                      /*}}}*/
  void write_vtu(const ptree& root, const _real& time) { /*{{{*/
    std::filesystem::path vtu_full_path = m_path;
    std::filesystem::path file_name =
      m_proj_name.string() + "_" + std::to_string(++m_count) + ".vtu";
    vtu_full_path.append(file_name.string());
    boost::property_tree::write_xml(
      vtu_full_path, root, std::locale(),
      boost::property_tree::xml_writer_make_settings<std::string>(' ', 2));
    ////////////////////////////////////////////////////////////////

    std::filesystem::path pvd_full_path = m_path;
    pvd_full_path.append(m_proj_name.string() + "_" + ".pvd");
    ptree& collection = m_pvd.get_child("VTKFile.Collection");
    ptree  DataSet;

    DataSet.put("<xmlattr>.timestep", time);
    DataSet.put("<xmlattr>.group", "");
    DataSet.put("<xmlattr>.part", "0");
    DataSet.put("<xmlattr>.file", file_name.string());
    collection.add_child("DataSet", DataSet);

    boost::property_tree::write_xml(
      pvd_full_path, m_pvd, std::locale(),
      boost::property_tree::xml_writer_make_settings<std::string>(' ', 2));
    return;
  }                    /*}}}*/
  ptree create_vtu() { /*{{{*/
    ptree root;
    ptree VTKFile;
    { /* populate VTKFile node */
      VTKFile.add("<xmlattr>.type", "UnstructuredGrid");
      VTKFile.add("<xmlattr>.version", "0.1");
      VTKFile.add("<xmlattr>.byte_order", "LittleEndian");
      ptree UnstructuredGrid;
      { /* populate UnstructuredGrid node */
        ptree Piece;
        { /* populate Piece node */
          ptree Points;
          ptree Cells;
          ptree PointData;
          ptree CellData;
          Piece.push_back(std::make_pair("Points", Points));
          Piece.push_back(std::make_pair("Cells", Cells));
          Piece.push_back(std::make_pair("PointData", PointData));
          Piece.push_back(std::make_pair("CellData", CellData));
        }
        UnstructuredGrid.push_back(std::make_pair("Piece", Piece));
      }
      VTKFile.push_back(std::make_pair("UnstructuredGrid", UnstructuredGrid));
    }
    root.push_back(std::make_pair("VTKFile", VTKFile));

    return root;
  } /*}}}*/

  ptree                 m_pvd;
  ptree                 m_gauss_pvd;
  std::filesystem::path m_path;
  std::filesystem::path m_proj_name;
  size_t                m_count;
  size_t                m_count2;
};

}  // namespace io
}  // namespace geoxfem

#endif /* VTK_WRITER_H */
